const express = require("express");
const router = express.Router();
const { spawn, spawnSync, exec } = require("child_process");
const path = require("path");
const verifyStatus = require("../middleware/verifyStatus");

router.get("/run", verifyStatus(), (req, res) => {
  let curSum = 0;
  let count = 0;

  // const child = spawn("python", ["test.py"], {
  //   cwd: path.join(process.cwd(), "app"),
  // });
  const child = exec(
    "bash -c 'source venv/bin/activate && python ./CV-RL-PR/RL/DQNv2/dqn_main.py --mode test'",
    {
      cwd: path.join(process.cwd(), "app"),
    }
  );
  child.stdout.on("data", (data) => {
    const msg = data.toString().trim();
    const strs = msg.split(" ");
    if (!strs[2]?.startsWith("global_step")) {
      return;
    }
    const reward = Number.parseInt(strs[5]?.split("=")[1]);

    curSum += reward;
    count++;

    const info = {
      x: strs[1],
      y: strs[5]?.split("=")[1],
      avg: curSum / count,
      image: strs[3]?.split("=")[1],
      cvmodel: strs[4]?.split("=")[1],
    };
    const seq = JSON.stringify(info);
    req.app.locals.wsTestLog?.send(msg);
    req.app.locals.wsTestChart?.send(seq);
    req.app.locals.wsTestResult?.send(seq);
  });

  child.stderr.on("data", (data) => {
    console.log(`${data}`);
    req.app.locals.wsTestLog.send(`${data}`);
  });

  child.on("close", (code) => {
    console.log(`子进程退出,退出码 ${code}`);
    req.app.locals.wsTestLog.send(`exit: ${code}`);
    req.app.locals.status = "WAITING";
  });
  req.app.locals.status = "TESTING";
  req.app.locals.testProcess = child;
  return res.json({ code: 0, msg: "训练任务启动成功" });
});

router.get("/stop", (req, res) => {
  const target = req.app.locals.testProcess?.pid + 1;
  spawnSync("kill", ["-9", `${target}`]);
  // req.app.locals.testProcess?.kill();
  return res.json({ code: 0, msg: "停止训练任务" });
});

router.get("/file", (req, res) => {
  if (req.query.type === "input") {
    const imgPath = path.join(
      process.cwd(),
      "app",
      "datasets",
      req.query.dataset,
      "val",
      "images",
      req.query.image
    );
    res.setHeader("Content-Type", "image/png");
    res.sendFile(imgPath);
  } else if (req.query.type === "output") {
    const imgPath = path.join(
      process.cwd(),
      "app",
      "datasets",
      req.query.dataset,
      "val",
      "res",
      req.query.cvmodel,
      "vis",
      req.query.image
    );
    res.setHeader("Content-Type", "image/png");
    res.sendFile(imgPath);
  } else {
    res.status(400).json({ code: 1, message: "type参数错误" });
  }
});

module.exports = router;
